#include "BWT.h"
#include <search.h>

//
// constants
//
#define BLOCK_SIZE 0x1000000

//
// prototypes
//
int __cdecl MyCompareProc( const void *elem1, const void *elem2 );

//
// global variables
//
UINT*    g_uiIndices;
char*    g_cBlock;
DWORD    g_cbBlock;

//
// Purpose:
//   Userdefined compare routine for the qsort() call
//
int __cdecl MyCompareProc( const void *elem1, const void *elem2 )
{
	UINT uiLen1 = g_cbBlock - *(UINT*)elem1;
	UINT uiLen2 = g_cbBlock - *(UINT*)elem2;

	// compare 2 strings being referenced by the 2 specified pointers to indices
	int res = memcmp(
		g_cBlock + *(UINT*)elem1,
		g_cBlock + *(UINT*)elem2,
		min(uiLen1,uiLen2) );

	if ( res == 0)
		return uiLen2 - uiLen1;
	else
		return res;
}

//
// Remarks:
//   The outbuffer output variable must be freed with the delete keyword when it
//   isn't needed anymore
//
BOOL BWT::Encode( IN PBYTE buffer, IN DWORD cb, OUT void* &outbuffer, OUT DWORD &cbOut )
{
	BOOL     bRet = FALSE;
	BYTE*    pOut = NULL;

	// pre-initialize output variables
	outbuffer  = NULL;
	cbOut      = 0;

	//
	// allocate some working memory
	//
	g_cBlock = new char[BLOCK_SIZE];
	if ( !g_cBlock )
		return FALSE; // ERR
	g_uiIndices = new UINT[BLOCK_SIZE + 1];
	if ( !g_uiIndices )
	{
		delete [] g_cBlock;
		return FALSE; // ERR
	}
	pOut = new BYTE[ ((cb/BLOCK_SIZE)+1)*12 + cb + 1 ];

	//
	// enter main loop
	// 
	// process the input buffer in BLOCK_SIZE pieces
	//
	DWORD cbDone  = 0;
	while( cbDone != cb )
	{
		// calc current block size
		g_cbBlock = min(cb - cbDone, BLOCK_SIZE);

		// stamp block size
		*(PDWORD)((DWORD)pOut + cbOut) = g_cbBlock+1; // blocksize + 1 termination symbol
		cbOut += 4;

		// fill indices chain with running integers starting from 0...c[0]=0;c[1]=1...
		for( UINT i = 0; i <= g_cbBlock; i++ )
			g_uiIndices[i] = i;
		// copy data -> g_cBlock
		memcpy( &g_cBlock[0], &buffer[cbDone], g_cbBlock ); 

		//
		// sort the indices pointing into the input buffer alphabetically
		//
		qsort( (void*)g_uiIndices, g_cbBlock+1, sizeof(UINT), MyCompareProc );

		//
		// stamp column L
		//
		// - special terminating end character (#256) = '?'
		// - save indices of first and last character of input for later reconstruction
		//
		UINT i1st, iLast;
		for( UINT u = 0; u <= g_cbBlock; u++,cbOut++ )
		{
			if ( g_uiIndices[u] == 1 )
				i1st = u;
			if ( g_uiIndices[u] == 0 )
			{
				iLast        = u;
				pOut[cbOut]  = '?';
			}
			else
				pOut[cbOut] = g_cBlock[ g_uiIndices[u] - 1 ];			
		}

		// stamp first/last index and adjust vars
		cbDone  += g_cbBlock;
		*(PDWORD)( (DWORD)pOut + cbOut) = i1st;
		cbOut   += 4;
		*(PDWORD)( (DWORD)pOut + cbOut) = iLast;
		cbOut   += 4;

	}

	outbuffer = pOut;

	//
	// tidy up
	//
	delete [] g_cBlock;
	delete [] g_uiIndices;

	bRet = TRUE; // return TRUE

	return bRet;
}

//
// Remarks:
//   The outbuffer output variable must be freed with the delete keyword when it
//   isn't needed anymore
//
//   NO arguments should be equivalent !
//
// Args:
//   cb  - the size of the orginal block before it was modifed by BWT::Encode
//
BOOL BWT::Decode( IN PBYTE buffer, IN DWORD cb, OUT void* &outbuffer, OUT DWORD &cbOut )
{
	BOOL bRet = FALSE;

	// pre-initialize output variables
	cbOut            = 0;
	outbuffer        = new BYTE[cb];
	if ( !outbuffer )
		return FALSE; // ERR
//	PBYTE outbuffer  = buffer; // use input buffer for output

	// allocate some working memory
	UINT *uiRunningTotal = new UINT[257];
	if ( !uiRunningTotal )
		return FALSE; // ERR
	UINT *uiCount = new UINT[257];
	if ( !uiCount )
	{
		delete [] uiRunningTotal;
		return FALSE; // ERR
	}
	UINT *T = new UINT[BLOCK_SIZE+1];
	if ( !T )
	{
		delete [] uiRunningTotal;
		delete [] uiCount;
		return FALSE; // ERR
	}

	//
	// in the near of the main loop
	//
	g_cBlock       = (CHAR*)buffer;
	DWORD cbDone   = 0;
	while( cbDone != cb )
	{
		// grab size of block
		g_cbBlock   = *(PDWORD)((DWORD)buffer + cbDone);
		cbDone     += 4;

		// grab first/last index being saved after the actual block
		UINT i1st   = *(PDWORD)((DWORD)buffer + cbDone + g_cbBlock + 0);
		UINT iLast  = *(PDWORD)((DWORD)buffer + cbDone + g_cbBlock + 4);

		// zero uiCount chain
		memset( (void*)uiCount, 0, 257*4 );

		//
		// Mark Nelson:
		//
		// To determine a character's position in the output string given
		// its position in the input string, we can use the knowledge about
		// the fact that the output string is sorted.  Each character 'c' will
		// show up in the output stream in in position i, where i is the sum
		// total of all characters in the input buffer that precede c in the
		// alphabet, plus the count of all occurences of 'c' previously in the
		// input stream.
		//
		// The first part of this code calculates the running totals for all
		// the characters in the alphabet.  That satisfies the first part of the
		// equation needed to determine where each 'c' will go in the output
		// stream.  Remember that the character pointed to by 'last' is a special
		// end-of-buffer character that is supposed to be larger than any char
		// in the alphabet.
		//
		for( UINT i = 0; i < g_cbBlock; i++ )
		{
			if ( i == iLast )
				uiCount[ 256 ]++;
			else
				uiCount[ buffer[i+cbDone] ]++;
		}

		DWORD dwSum = 0;
		for ( UINT u = 0; u < 257; u++ )
		{
			uiRunningTotal[u]   = dwSum;
			dwSum              += uiCount[u];
			uiCount[u]          = 0;
		}

		//
		// Mark Nelson:
		//
		// Now that the RunningTotal[] array is filled in, I have half the
		// information needed to position each 'c' in the input buffer.  The
		// next piece of information is simply the number of characters 'c'
		// that appear before this 'c' in the input stream.  I keep track of
		// that informatin in the Counts[] array as I go.  By adding those
		// two number together, I get the destination of each character in
		// the input buffer, and that allows me to fill in all the positions
		// in the transformation vector, T[].
		//
		for ( UINT z = 0; z < g_cbBlock; z++ )
		{
			UINT index;
			if ( z == iLast )
                index = 256; // index = pseudo-character
			else
				index = buffer[cbDone+z];

            T[ uiCount[index] + uiRunningTotal[index] ] = z;
            uiCount[ index ]++;
		}

		// finally use the transformation vector T to build the output
		UINT x = i1st;
		for ( UINT y = 0; y < g_cbBlock-1; y++ )
		{
			((PBYTE)outbuffer)[cbOut++]  = buffer[x+cbDone];
			x                            = T[x];
		}

		// adjust vars; process next block
		cbDone     += g_cbBlock + 8;
		g_cBlock    = (CHAR*)(buffer + cbDone);
	}

	//
	// tidy up
	//
	delete [] uiRunningTotal;
	delete [] uiCount;
	delete [] T;

	bRet = TRUE; // return TRUE

	return bRet;
}